#include "config.h"
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <errno.h>
#include <netdb.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <arpa/inet.h>
#include <libgen.h>
#include "cJSON/cJSON.h"
#include "io.h"
#include "misc.h"
#include "physmem.h"

#define NEED_FREE_SERVER_INFO_ADDR      1
#define NEED_FREE_SERVER_INFO_PASSWD    2
#define NEED_FREE_SERVER_INFO_METHOD    4

static void usage(const char *cmd, const char *prog) {
    fprintf(stdout,
"Usage: %s [OPTION]...\n"
"Options:\n"
"  -c | --config <pathname>            path to config file. [%s]\n"
"  -s | --server_addr <server addr>    server address."
"\n"
"  -p | --server_port <server port>    server port. [%s]\n"
"  -b | --local_bind <bind addr>       local binding address. [%s]\n"
"  -l | --local_listen <listen port>   local listen port. [%s]\n"
"  -k | --password <password>          server password.\n"
"  -m | --method <method>              encryption method. [%s]\n"
"  -t | --timeout <timeout>            socket timeout in seconds. [%ds]\n"
"     | --connect_timeout <timeout>    timeout of connect to server. [%ds]\n"
"     | --accept_timeout <timeout>     the timeout of new connection after accept. [%ds]\n"
"  -e | --expire <timeout>             dns cache expire timeout. [%ds]\n"
"\n"
"  -h | --help                         show this help message and exit.\n"
"       --daemon                       daemon mode.\n"
"       --pid-file <pid file>          ["DEF_PID_FILE"]\n"
"       --version                      show version information.\n"
            , cmd, DEF_CFG_FILE, DEF_SRV_PORT, DEF_BIND_ADDR, DEF_LISTEN_PORT,
            DEF_CRYPTO, DEF_SOCKET_TIMEOUT, DEF_CONNECT_TIMEOUT,
            DEF_ACCEPT_TIMEOUT, DEF_CACHE_EXPIRE, prog);
    exit(EXIT_FAILURE);
}

static char *options = "b:c:e:hk:l:m:p:s:t:";
static struct option longOptions[] = {
        { "help", no_argument, 0, 'h' },
        { "config", required_argument, 0, 'c' },
        { "server_addr", required_argument, 0, 's' },
        { "server_port", required_argument, 0, 'p' },
        { "local_addr", required_argument, 0, 'b' },
        { "local_port", required_argument, 0, 'l' },
        { "password", required_argument, 0, 'k' },
        { "method", required_argument, 0, 'm' },
        { "timeout", required_argument, 0, 't' },
        { "expire", required_argument, 0, 'e' },
        { "connect_timeout", required_argument, 0, 0 },
        { "accept_timeout", required_argument, 0, 0 },
        { "pid-file", required_argument, 0, 0 },
        { "daemon", no_argument, 0, 0 },
        { "version", no_argument, 0, 0 },
        { NULL, 0, NULL, 0 }
};

struct srv_info_from_cmd {
    const char *service;
    const char *passwd;
    const char *crypt_method;
};

static in_port_t service2port(const char *service, const char *proto) {
    in_port_t port;
    char *c;

    port = strtoul (service, &c, 10);
    if (*c) {
        struct servent *servent = getservbyname(service, proto);
        if (servent)
            port = servent->s_port;
        else
            port = 0;
    }

    return htons(port);
}

static char *istrdup(const char *src) {
    char *dup = strdup(src);
    if (!dup) {
        fprintf(stderr, "Not enough memory\n");
        exit(EXIT_FAILURE);
    }
    return dup;
}

static void *irealloc(void *ptr, size_t size) {
    void *ret = realloc(ptr, size);
    if (!ret) {
        fprintf(stderr, "Not enough memory\n");
        exit(EXIT_FAILURE);
    }
    return ret;
}

static void get_servers_old_fmt(struct cfg *cfg, cJSON *root) {
    in_port_t port;
    char *passwd, *method;
    cJSON *item;

    if ((item = cJSON_GetObjectItem(root, "server_port"))) {
        if (item->type == cJSON_Number) {
            port = htons(item->valueint);
        } else if (item->type == cJSON_String) {
            port = service2port(item->valuestring, "tcp");
        } else {
            return;
        }
    } else {
        return;
    }

    if ((item = cJSON_GetObjectItem(root, "password")) && item->type == cJSON_String)
        passwd = istrdup(item->valuestring);
    else
        return;

    if ((item = cJSON_GetObjectItem(root, "method")) && item->type == cJSON_String)
        method = istrdup(item->valuestring);
    else {
        free(passwd);
        return;
    }

    if ((item = cJSON_GetObjectItem(root, "server"))) {
        int i = cfg->server_cnt;
        if (item->type == cJSON_Array) {
            int n = cJSON_GetArraySize(item);
            if (!n) {
                free(passwd);
                free(method);
                return;
            }
            cfg->server_infos = irealloc(cfg->server_infos,
                    sizeof(*cfg->server_infos) * (cfg->server_cnt + n));
            for (item = item->child; item; ++i, item = item->next) {
                cfg->server_infos[i].addr = istrdup(item->valuestring);
                cfg->server_infos[i].port = port;
                cfg->server_infos[i].passwd = passwd;
                cfg->server_infos[i].crypt_method = method;
                cfg->server_infos[i].flags = NEED_FREE_SERVER_INFO_ADDR;
            }
            cfg->server_infos[cfg->server_cnt].flags |=
                    NEED_FREE_SERVER_INFO_PASSWD | NEED_FREE_SERVER_INFO_METHOD;
            cfg->server_cnt = i;
        } else if (item->type == cJSON_String) {
            cfg->server_infos = irealloc(cfg->server_infos,
                    sizeof(*cfg->server_infos) * (cfg->server_cnt + 1));
            cfg->server_infos[i].addr = istrdup(item->valuestring);
            cfg->server_infos[i].port = port;
            cfg->server_infos[i].passwd = passwd;
            cfg->server_infos[i].crypt_method = method;
            cfg->server_infos[cfg->server_cnt].flags =
                    NEED_FREE_SERVER_INFO_ADDR | NEED_FREE_SERVER_INFO_PASSWD
                            | NEED_FREE_SERVER_INFO_METHOD;
            cfg->server_cnt += 1;
        } else {
            free(passwd);
            free(method);
        }
    } else {
        free(passwd);
        free(method);
    }
}

static void get_server(struct cfg *cfg, cJSON *root) {
    in_port_t port;
    char *passwd, *method;
    cJSON *item;

    if (!(item = cJSON_GetObjectItem(root, "port")))
        return;
    if (item->type == cJSON_Number) {
        port = item->valueint;
    } else if (item->type == cJSON_String) {
        port = service2port(item->valuestring, "tcp");
    } else {
        return;
    }

    if ((item = cJSON_GetObjectItem(root, "passwd")) && item->type == cJSON_String)
        passwd = istrdup(item->valuestring);
    else
        return;

    if ((item = cJSON_GetObjectItem(root, "crypt_method")) && item->type == cJSON_String)
        method = istrdup(item->valuestring);
    else {
        free(passwd);
        return;
    }

    if ((item = cJSON_GetObjectItem(root, "addr"))) {
        int i = cfg->server_cnt;
        if (item->type == cJSON_Array) {
            int n = cJSON_GetArraySize(root);
            if (!n) {
                free(passwd);
                free(method);
                return;
            }
            cfg->server_infos = irealloc(cfg->server_infos,
                    sizeof(*cfg->server_infos) * (cfg->server_cnt + n));
            for (item = item->child; item; ++i, item = item->next) {
                cfg->server_infos[i].addr = istrdup(item->valuestring);
                cfg->server_infos[i].port = port;
                cfg->server_infos[i].passwd = passwd;
                cfg->server_infos[i].crypt_method = method;
                cfg->server_infos[i].flags = NEED_FREE_SERVER_INFO_ADDR;
            }
            cfg->server_infos[cfg->server_cnt].flags |=
                    NEED_FREE_SERVER_INFO_PASSWD | NEED_FREE_SERVER_INFO_METHOD;
            cfg->server_cnt = i;
        } else if (item->type == cJSON_String) {
            cfg->server_infos = irealloc(cfg->server_infos,
                    sizeof(*cfg->server_infos) * (cfg->server_cnt + 1));
            cfg->server_infos[i].addr = istrdup(item->valuestring);
            cfg->server_infos[i].port = port;
            cfg->server_infos[i].passwd = passwd;
            cfg->server_infos[i].crypt_method = method;
            cfg->server_infos[cfg->server_cnt].flags =
                    NEED_FREE_SERVER_INFO_ADDR | NEED_FREE_SERVER_INFO_PASSWD
                            | NEED_FREE_SERVER_INFO_METHOD;
            cfg->server_cnt += 1;
        } else {
            free(passwd);
            free(method);
        }
    } else {
        free(passwd);
        free(method);
    }
}

static void get_servers(struct cfg *cfg, cJSON *root) {
    if (!(root = cJSON_GetObjectItem(root, "servers")))
        return;

    if (root->type == cJSON_Array) {
        cJSON *c;
        for (c = root->child; c; c = c->next)
            get_server(cfg, c);
    } else if (root->type == cJSON_Object) {
        get_server(cfg, root);
    }
}

static void config_file_load(struct cfg *cfg, const char *file, int errexit,
        struct cfg *from_file) {
    int fd = -1;
    char *data = NULL;

    memset(from_file, 0, sizeof(*from_file));

    if ((fd = open(file, O_RDONLY)) < 0) {
        if (errexit)
            fprintf(stderr, "Cannot open config file \"%s\": %s\n", file, strerror(errno));
        goto ERR;
    }
    struct stat stat;
    if (fstat(fd, &stat) < 0) {
        fprintf(stderr, "Cannot get file info for \"%s\": %s\n", file, strerror(errno));
        goto ERR;
    }
    if (!(data = malloc(stat.st_size + 1))) {
        fprintf(stderr, "Not enough memory\n");
        exit(EXIT_FAILURE);
    }
    if (full_read(fd, data, stat.st_size) != stat.st_size) {
        fprintf(stderr, "Read file \"%s\" error: %s\n", file, strerror(errno));
        goto ERR;
    }
    data[stat.st_size] = '\0';
    close(fd);
    fd = -1;

    struct cfg c = *cfg;

    cJSON *root, *item;
    if (!(root = cJSON_Parse(data))) {
        fprintf(stderr, "Error before: [%s]\n", cJSON_GetErrorPtr());
        goto ERR;
    }

    get_servers_old_fmt(from_file, root);
    get_servers(from_file, root);
    if (from_file->server_infos) {
        c.server_infos = from_file->server_infos;
        c.server_cnt = from_file->server_cnt;
    }

    if ((item = cJSON_GetObjectItem(root, "local_address")) && item->type == cJSON_String) {
        from_file->local_bind = istrdup(item->valuestring);
        c.local_bind = from_file->local_bind;
    }

    if ((item = cJSON_GetObjectItem(root, "local_port"))) {
        if (item->type == cJSON_Number) {
            char service[32];
            snprintf(service, sizeof(service), "%u", item->valueint);
            from_file->local_listen = istrdup(service);
        } else if (item->type == cJSON_String) {
            from_file->local_listen = istrdup(item->valuestring);
        }
        c.local_listen = from_file->local_listen;
    }

    if ((item = cJSON_GetObjectItem(root, "timeout")) && item->type == cJSON_Number) {
        from_file->socket_timeout = item->valueint;
        c.socket_timeout = from_file->socket_timeout;
    }

    if ((item = cJSON_GetObjectItem(root, "connect_timeout")) && item->type == cJSON_Number) {
        from_file->connect_timeout = item->valueint;
        c.connect_timeout = from_file->connect_timeout;
    }

    if ((item = cJSON_GetObjectItem(root, "accept_timeout")) && item->type == cJSON_Number) {
        from_file->accept_timeout = item->valueint;
        c.accept_timeout = from_file->accept_timeout;
    }


    if ((item = cJSON_GetObjectItem(root, "dns_cache_expire")) && item->type == cJSON_Number) {
        from_file->dns_cache_expire = item->valueint;
        c.dns_cache_expire = from_file->dns_cache_expire;
    }

    if ((item = cJSON_GetObjectItem(root, "daemon"))) {
        if (item->type == cJSON_False)
            from_file->daemon = 0;
        else if (item->type == cJSON_True)
            from_file->daemon = 1;
        else if (item->type == cJSON_Number)
            from_file->daemon = item->valueint;
        c.daemon = from_file->daemon;
    }

    if ((item = cJSON_GetObjectItem(root, "pid-file")) && item->type == cJSON_String) {
        from_file->pidfile = istrdup(item->valuestring);
        c.pidfile = from_file->pidfile;
    }

    cJSON_Delete(root);

    free(data);

    *cfg = c;
    return;
ERR:
    if (fd >= 0) close(fd);
    if (data) free(data);
    if (errexit) exit(EXIT_FAILURE);
}

static void server_info_change_passwd(struct cfg *cfg, const char *value) {
    int i;
    for (i = 0; i < cfg->server_cnt; ++i) {
        if (cfg->server_infos[i].flags & NEED_FREE_SERVER_INFO_PASSWD) {
            free((void *)cfg->server_infos[i].passwd);
            cfg->server_infos[i].flags &= ~NEED_FREE_SERVER_INFO_PASSWD;
        }
        cfg->server_infos[i].passwd = value;
    }
}

static void server_info_change_crypt_method(struct cfg *cfg, const char *value) {
    int i;
    for (i = 0; i < cfg->server_cnt; ++i) {
        if (cfg->server_infos[i].flags & NEED_FREE_SERVER_INFO_METHOD) {
            free((void *)cfg->server_infos[i].passwd);
            cfg->server_infos[i].flags &= ~NEED_FREE_SERVER_INFO_METHOD;
        }
        cfg->server_infos[i].passwd = value;
    }
}

static void server_info_change_port(struct cfg *cfg, const char *value) {
    int i;
    in_port_t port = service2port(value, "tcp");
    if (!port) {
        fprintf(stderr, "Invalid service specified: \"%s\"\n", value);
        exit(EXIT_FAILURE);
    }

    for (i = 0; i < cfg->server_cnt; ++i)
        cfg->server_infos[i].port = port;
}

static void free_server_infos(struct cfg *cfg) {
    if (cfg->server_cnt) {
        int i;
        for (i = 0; i < cfg->server_cnt; ++i) {
            if (cfg->server_infos[i].flags & NEED_FREE_SERVER_INFO_ADDR)
                free((void *)cfg->server_infos[i].addr);
            if (cfg->server_infos[i].flags & NEED_FREE_SERVER_INFO_PASSWD)
                free((void *)cfg->server_infos[i].passwd);
            if (cfg->server_infos[i].flags & NEED_FREE_SERVER_INFO_METHOD)
                free((void *)cfg->server_infos[i].crypt_method);
        }
        free(cfg->server_infos);
        cfg->server_infos = NULL;
        cfg->server_cnt = 0;
    }
}

static void server_info_change_addr(struct cfg *cfg, char *value,
        struct srv_info_from_cmd *s) {
    in_port_t port;
    const char *passwd, *method;
    if (cfg->server_cnt) {
        port = cfg->server_infos[0].port;
        passwd = istrdup(cfg->server_infos[0].passwd);
        method = istrdup(cfg->server_infos[0].crypt_method);
    } else {
        port = service2port(s->service, "tcp");
        passwd = s->passwd;
        method = s->crypt_method;
    }
    free_server_infos(cfg);

    char *one;
    char *saveptr;
    int i = 0;

    for (one = strtok_r(value, ",", &saveptr); one;
            one = strtok_r(NULL, ",", &saveptr), ++i) {
        cfg->server_infos = irealloc(cfg->server_infos, sizeof(*cfg->server_infos) * (i + 1));
        cfg->server_infos[i].addr = one;
        cfg->server_infos[i].port = port;
        cfg->server_infos[i].passwd = passwd;
        cfg->server_infos[i].crypt_method = method;
        cfg->server_infos[i].flags = 0;
    }
    if (i)
        cfg->server_infos[0].flags |= NEED_FREE_SERVER_INFO_PASSWD
                | NEED_FREE_SERVER_INFO_METHOD;
    cfg->server_cnt = i;
}

static unsigned long optargtoul(const char * argv0, const char *prog_name,
        const char *optarg, const char *errstr)
{
    unsigned long tmpint;
    char *tmpcptr;
    tmpint = strtoul(optarg, &tmpcptr, 10);
    if (*tmpcptr) {
        fprintf(stderr, "%s: \"%s\"\n", errstr, optarg);
        usage(argv0, prog_name);
        return 0;
    }
    return tmpint;
}

static void igetopt(struct cfg *cfg, int argc, char *argv[]) {
    int c;
    int optindb = optind;
    struct cfg from_file;
    struct srv_info_from_cmd srvinfocmd = {
            .service = DEF_SRV_PORT,
            .passwd = NULL,
            .crypt_method = DEF_CRYPTO,
    };

    /* Parse configure file at first */
    while (1) {
        if ((c = getopt_long(argc, argv, options, longOptions, NULL)) == -1)
            break;
        if (c == 'c') {
            cfg->config_file = optarg;
            break;
        } else if (c == 'h' || c == '?') {
            usage(argv[0], cfg->prog_name);
        }
    }
    if (cfg->config_file)
        config_file_load(cfg, cfg->config_file, 1, &from_file);
    else
        config_file_load(cfg, DEF_CFG_FILE, 0, &from_file);

    optind = optindb;
    while (1) {
        int opt_idx;
        if ((c = getopt_long(argc, argv, options, longOptions, &opt_idx)) == -1)
            break;
        switch (c) {
        case 0:
            if (strcmp("connect_timeout", longOptions[opt_idx].name) == 0) {
                cfg->connect_timeout = optargtoul(argv[0], cfg->prog_name, optarg, "Invalid connect timeout");
            } else if (strcmp("accept_timeout", longOptions[opt_idx].name) == 0) {
                cfg->accept_timeout = optargtoul(argv[0], cfg->prog_name, optarg, "Invalid accept timeout");
            } else if (strcmp("pid-file", longOptions[opt_idx].name) == 0) {
                cfg->pidfile = optarg;
            } else if (strcmp("daemon", longOptions[opt_idx].name) == 0) {
                cfg->daemon = 1;
            }
            break;
        case 'b':
            if (from_file.local_bind)
                free((void *)cfg->local_bind);
            cfg->local_bind = optarg;
            break;
        case 'e':
            cfg->dns_cache_expire = optargtoul(argv[0], cfg->prog_name, optarg, "Invalid expire timeout");
            break;
        case 'k':
            srvinfocmd.passwd = optarg;
            server_info_change_passwd(cfg, optarg);
            break;
        case 'l':
            if (from_file.local_listen)
                free((void *)cfg->local_listen);
            cfg->local_listen = optarg;
            break;
        case 'm':
            srvinfocmd.crypt_method = optarg;
            server_info_change_crypt_method(cfg, optarg);
            break;
        case 'p':
            srvinfocmd.service = optarg;
            server_info_change_port(cfg, optarg);
            break;
        case 's':
            server_info_change_addr(cfg, optarg, &srvinfocmd);
            break;
        case 't':
            cfg->socket_timeout = optargtoul(argv[0], cfg->prog_name, optarg, "Invalid socket timeout");
            break;
        default:
            break;
        }
    }

    if (cfg->server_cnt < 1) {
        fprintf(stderr, "There is no server specified\n");
        exit(EXIT_FAILURE);
    }

    if (!cfg->server_infos[0].passwd) {
        fprintf(stderr, "Password not specified\n");
        usage(argv[0], cfg->prog_name);
    }
}

static void get_open_max(struct cfg *cfg) {
    struct rlimit rl;
    /* Get maximum number of file descriptors. */
    if (getrlimit(RLIMIT_NOFILE, &rl) < 0) {
        fprintf(stderr, "%s: can't get file limit: getrlimit error: %s\n",
                cfg->prog_name, strerror(errno));
        exit(EXIT_FAILURE);
    }
    cfg->open_max = rl.rlim_cur;
    if (rl.rlim_max == RLIM_INFINITY)
        rl.rlim_max = OPEN_MAX_DEF;
    rl.rlim_cur = rl.rlim_max;
    if (setrlimit(RLIMIT_NOFILE, &rl) == 0)
        cfg->open_max = rl.rlim_cur;
}

static void set_buf_size(struct cfg *cfg) {
    size_t available_mem = physmem_available();
    if (available_mem) {
        cfg->buf_size = available_mem / cfg->open_max;
        cfg->buf_size = rounddown_pow_of_2(cfg->buf_size);
        if (cfg->buf_size > BUF_SIZE_MAX)
            cfg->buf_size = BUF_SIZE_MAX;
        else if (cfg->buf_size < BUF_SIZE_MIN)
            cfg->buf_size = BUF_SIZE_MIN;
    } else
        cfg->buf_size = BUF_SIZE_DEF;
}

static void make_rand_seed(struct cfg *cfg) {
    unsigned seed;
    int fd = open("/dev/random", O_RDONLY);
    if (fd < 0 || full_read(fd, &seed, sizeof(seed)) != sizeof(seed)) {
        /* faild, we use time() as rand seed */
        seed = time(NULL);
    }
    if (fd >= 0)
        close(fd);

    cfg->rand_seed = seed;
}

void config(struct cfg *cfg, int argc, char *argv[]) {
    cfg->prog_name = basename(argv[0]);
    igetopt(cfg, argc, argv);

    int i;
    for (i = 0; i < cfg->server_cnt; ++i)
        cfg->server_infos[i].weight = cfg->server_initial_weight;
    cfg->server_available_idx = 0;

    get_open_max(cfg);
    set_buf_size(cfg);
    make_rand_seed(cfg);
}
